'''
读取csv格式文件，提取前8列, 以csv格式保存成新的文件。采用命令行参数传入输入、输出的文件名
'''
import pandas as pd
import sys
import argparse

# 创建ArgumentParser对象
parser = argparse.ArgumentParser()

# 添加两个必需的string类型命令行参数
parser.add_argument('input_file', type=str, help='Enter the input file path')
parser.add_argument('output_file', type=str, help='Enter the output file path')

parser.add_argument('--truncate', type=int, default=0, 
                    help='truncate rows for row index > truncate, default not truncating')

# 解析命令行参数
args = parser.parse_args()

data = pd.read_csv(args.input_file, sep=' ',header=None, comment='#')

# truncate rows for large file
if args.truncate:
    max_row = args.truncate
    print(f'truncate rows for row index >= {max_row}.')
    data = data.iloc[:max_row, :]

# get the first 8 columns
# new_data = data.iloc[:, :8]

new_data = data.iloc[:, :3]
# 给数据列增加header
new_header = ['timestamp', 'x', 'y']
new_data.columns = new_header

# add columns
new_data['z'] = 0.0
new_data['qx'] = 0.0
new_data['qy'] = 0.0
new_data['qz'] = 0.0
new_data['qw'] = 1.0

print(new_data.head())

print("save data to file: ", args.output_file)
new_data.to_csv(args.output_file, sep=' ', header=False, index=False)

